from DEAD.AutoDecoder.Evaluation.Decoding import *
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from matplotlib.ticker import MultipleLocator, AutoMinorLocator
import numpy as np
import umap  # 需要安装 umap-learn: pip install umap-learn
import joblib  # 需要安装：pip install joblib
from DDPM.Config import decode_result_save_path
from DDPM.Evaluation.Control import *

def label_map(arr, intervals):
    """
    对输入的numpy数组arr的每个元素进行判断，根据预定义的区间长度将其映射到相应的值。

    :param arr: 输入的numpy数组
    :param intervals: 区间长度列表，格式为 [length1, length2, ...]
    :return: 映射后的numpy数组
    """
    # 初始化结果数组，初始值为原始数组的值
    mapped_arr = np.copy(arr)
    
    # 初始化区间起点
    start = arr.min()
    
    # 遍历每个区间
    for i, length in enumerate(intervals):
        # 计算区间终点
        end = start + length
        # 使用numpy的where函数进行条件判断和赋值
        mapped_arr = np.where((arr >= start) & (arr < end), i, mapped_arr)
        # 更新区间起点
        start = end
    
    return mapped_arr


def train_umap(feature_vectors, umap_save_path, std):
    # 初始化 UMAP
    reducer = umap.UMAP(
        n_components=2,      # 降维到2D
        n_neighbors=1000,      # 控制局部/全局结构平衡（类似 t-SNE 的 perplexity）
        min_dist=0.5,        # 控制点之间的最小距离（越小聚类越紧凑）
        metric='euclidean',  # 距离度量方式
        random_state=42      # 随机种子
    )

    # 训练 UMAP 并获取初始嵌入
    print("开始训练 UMAP")
    embedding = reducer.fit_transform(feature_vectors/std)

    # 创建一个字典来保存 UMAP 模型和 std 值
    umap_dict = {
        'reducer': reducer,
        'std': std
    }

    # 训练UMAP后保存
    joblib.dump(umap_dict, umap_save_path)
    print(f"UMAP模型已保存至 {umap_save_path}")
    return embedding, reducer

def load_umap(umap_save_path):
    # 加载现有模型
    model_dict = joblib.load(umap_save_path)
    reducer = model_dict['reducer']
    std = model_dict['std']
    return reducer, std

def reduce_feature_vectors(feature_vectors, std, reducer):
    feature_vectors = feature_vectors / std
    embedding = reducer.transform(feature_vectors)
    return embedding

def load_feature_vectors(npz_path,npz_name):
    npz_file = np.load(f"{npz_path}/{npz_name}.npz")
    lvs = npz_file["lv"]
    gvs = npz_file["gv"]
    feature_vectors = np.concatenate((lvs, gvs), axis=1)
    return feature_vectors

def load_feature_vectors_from_dead():
    from DEAD.AutoDecoder.Config import hidden_size, depth, data_num, latent_size, model_load_path
    model_load_path   # 加载模型和数据
    dead = setup_dead_model(hidden_size, depth, data_num, latent_size, model_load_path)
    latent_vectors = dead.latent_vectors.detach().cpu().numpy()
    geometic_vectors = dead.geometry_vectors.detach().cpu().numpy()
    std= latent_vectors.std()
    feature_vectors = np.concatenate((latent_vectors, geometic_vectors), axis=1)

    # 推算药型标签
    #参见doc文件夹下数据集的自由度
    intervals = [2, 5, 7, 7, 7, 6,7,8,11,12,6,6,6,6,12,9,\
             12,14,17,15,16,8,15,13,17,20,16,17,22,20,\
             15,11,17,15,16,20,7,5,13,8,4,8]
    num_of_dofs=np.sum(intervals)
    indices = np.arange(len(feature_vectors))
    remainders = (indices % (num_of_dofs*10))/10
    labels=label_map(remainders, intervals)

    return feature_vectors, labels, model_load_path, std

def show_u_map_clusterd(save_directory, old_embedding, new_embedding, clustered_embedding = None, plot_name="distribution_clusterd.png"):
    plt.rcParams['font.family'] = 'Times New Roman'
    plt.rcParams['font.size'] = 20  # 设置默认字体大小

    # 可视化最终结果
    plt.figure(figsize=(10, 8))

    plt.scatter(
        old_embedding[:, 0],
        old_embedding[:, 1],
        c=[(0.7,0.7,0.7)],  # 指定颜色为灰色
        alpha=0.1,
        s=10
    )

    # 叠加绘制 new_embedding，颜色为蓝色
    plt.scatter(
        new_embedding[:, 0],
        new_embedding[:, 1],
        c='DarkBlue',  # 指定颜色为蓝色
        alpha=0.5,
        s=20,
        edgecolors='White',  # 设置边框颜色为白色
        linewidths=0.5,  # 设置边框宽度为1
        label='New Embedding'
    )

    if clustered_embedding is not None:
        # 叠加绘制 clustered_embedding，颜色为红色
        plt.scatter(
            clustered_embedding[:, 0],
            clustered_embedding[:, 1],
            c='Red',  # 指定颜色为红色
            alpha=1,
            s=20,
            edgecolors='White',  # 设置边框颜色为白色
            linewidths=0.5,  # 设置边框宽度为1
            label='Clustered Embedding'
        )

    # 设置主刻度和子刻度
    ax = plt.gca()
    ax.xaxis.set_minor_locator(AutoMinorLocator(4))  # 每个主刻度之间4个子刻度
    ax.yaxis.set_minor_locator(AutoMinorLocator(4))  # 每个主刻度之间4个子刻度
    
    plt.tick_params(axis='both', size=5, width=2, direction='in', labelsize=20)
    plt.tick_params(axis='both', which='minor', size=3, width=1, direction='in')

    # 设置横纵坐标比例相同
    plt.axis('equal')

    #plt.title("UMAP Projection of Latent Space")
    plt.xlabel("UMAP Dimension 1")
    plt.ylabel("UMAP Dimension 2")

    plt.savefig(f"{save_directory}/{plot_name}", dpi=600, transparent=False, format='png')
    plt.pause(3)
    plt.close()

def adjust_saturation(rgb_colors, saturation_factor):
    from matplotlib.colors import rgb_to_hsv, hsv_to_rgb
    has_alpha = rgb_colors.shape[1] == 4
    alphas = rgb_colors[:, 3] if has_alpha else None
    rgb = rgb_colors[:, :3]
    hsv = rgb_to_hsv(rgb[np.newaxis, :, :])[0]
    hsv[:, 1] = np.clip(hsv[:, 1] * saturation_factor, 0, 1)
    adjusted_rgb = hsv_to_rgb(hsv[np.newaxis, :, :])[0]
    if has_alpha:
        adjusted_colors = np.hstack((adjusted_rgb, alphas.reshape(-1, 1)))
    else:
        adjusted_colors = adjusted_rgb
    return adjusted_colors

def show_u_map(save_directory, labels, old_embedding):
    plt.rcParams['font.family'] = 'Times New Roman'
    plt.rcParams['font.size'] = 20  # 设置默认字体大小

    # 可视化最终结果
    plt.figure(figsize=(12, 8))

    # 设置参数
    cmap_name = 'hsv'
    n_colors = 42
    saturation_factor = 0.75  # 调整这里改变饱和度

    # 创建调色板
    cmap = plt.get_cmap(cmap_name)
    colors = cmap(np.linspace(0, 1, n_colors))
    adjusted_colors = adjust_saturation(colors, saturation_factor)
    custom_cmap = ListedColormap(adjusted_colors)

    plt.scatter(
        old_embedding[:, 0],
        old_embedding[:, 1],
        c=labels,  # 药型标签
        cmap=custom_cmap,
        alpha=0.5,
        s=10
    )
    cbar=plt.colorbar(label='Sample Index')
    cbar.solids.set(alpha=1)
    # 设置 colorbar 的主刻度和子刻度
    cbar.ax.yaxis.set_minor_locator(AutoMinorLocator(4))  # 每个主刻度之间4个子刻度

    # 设置主刻度和子刻度
    ax = plt.gca()
    ax.xaxis.set_minor_locator(AutoMinorLocator(4))  # 每个主刻度之间4个子刻度
    ax.yaxis.set_minor_locator(AutoMinorLocator(4))  # 每个主刻度之间4个子刻度
    
    plt.tick_params(axis='both', size=5, width=2, direction='in', labelsize=20)
    plt.tick_params(axis='both', which='minor', size=3, width=1, direction='in')

    # 设置横纵坐标比例相同
    plt.axis('equal')

    #plt.title("UMAP Projection of Latent Space")
    plt.xlabel("UMAP Dimension 1")
    plt.ylabel("UMAP Dimension 2")

    plt.savefig(f"{save_directory}/distribution.png", dpi=600, transparent=False, format='png')
    plt.pause(3)
    plt.close()

def start_umap(control_params):
    print(f"##### Start umap of {control_params['project_name']} #####")
    feature_vectors, labels, model_load_path, std=load_feature_vectors_from_dead()
    umap_save_path = f"{model_load_path}/umap_model.joblib"
    if not os.path.exists(umap_save_path):
        print("UMAP模型不存在, 需要训练UMAP模型")
        embedding, reducer = train_umap(feature_vectors, umap_save_path, std)
    else:
        print(f"UMAP模型已存在: {umap_save_path},加载现有模型")
        reducer, std= load_umap(umap_save_path)
        embedding = reduce_feature_vectors(feature_vectors, std, reducer)
        project_path=f"{decode_result_save_path}/{control_params['project_name']}"
        new_feature_vectors=load_feature_vectors(project_path,"Samples")
        new_embedding = reduce_feature_vectors(new_feature_vectors, std, reducer)

        # 绘制DEAD数据分布
        show_u_map(project_path, labels, embedding)

        # 绘制生成结果数据分布
        show_u_map_clusterd(project_path, embedding, new_embedding, None,"distribution_samples.png")

if __name__ == '__main__':
    control_params=get_control_params(question_name="K阶跃装药",
                postfix="_epoch50000",
                cluster_num=5,
                guid_w=0.5)
        
    start_umap(control_params)